#![windows_subsystem = "windows"]

use tray_icon::{
    menu::{Menu, MenuEvent, MenuItem, CheckMenuItem},
    TrayIcon, TrayIconBuilder, icon::Icon
};
use std::path::PathBuf;
use serde::{Serialize, Deserialize};
use std::sync::Arc;
use std::time::Duration;
use cpal::traits::{DeviceTrait, HostTrait, StreamTrait};
use parking_lot::Mutex;
use ringbuf::{HeapRb, Producer, Consumer};
use chrono::Local;
use winit::event_loop::EventLoopBuilder;
use log::{info, error};
use crossbeam_channel;

// 在文件开头添加常量定义
const MENU_ID_SAVE_5MIN: u32 = 1000;
const MENU_ID_SAVE_10MIN: u32 = 1001;
const MENU_ID_SAVE_30MIN: u32 = 1002;
const MENU_ID_OPEN_STARTUP: u32 = 1003;
const MENU_ID_ENHANCE_VOICE: u32 = 1004;
const MENU_ID_QUIT: u32 = 1005;

// 错误处理
#[derive(thiserror::Error, Debug)]
pub enum AppError {
    #[error("音频设备错误: {0}")]
    AudioError(String),

    #[error("音频流错误: {0}")]
    StreamError(String),

    #[error("WAV错误: {0}")]
    WavError(#[from] hound::Error),

    #[error("IO错误: {0}")]
    IoError(#[from] std::io::Error),

    #[error("配置错误: {0}")]
    ConfigError(#[from] serde_json::Error),

    #[error("图标错误: {0}")]
    IconError(String),

    #[error("菜单错误: {0}")]
    MenuError(String),

    #[error("系统错误: {0}")]
    SystemError(String),
}

// From trait 实现
impl From<cpal::BuildStreamError> for AppError {
    fn from(err: cpal::BuildStreamError) -> Self {
        Self::AudioError(err.to_string())
    }
}

impl From<cpal::PlayStreamError> for AppError {
    fn from(err: cpal::PlayStreamError) -> Self {
        Self::StreamError(err.to_string())
    }
}

impl From<cpal::DefaultStreamConfigError> for AppError {
    fn from(err: cpal::DefaultStreamConfigError) -> Self {
        Self::AudioError(err.to_string())
    }
}

impl From<String> for AppError {
    fn from(s: String) -> Self {
        Self::SystemError(s)
    }
}

impl From<&str> for AppError {
    fn from(s: &str) -> Self {
        Self::SystemError(s.to_string())
    }
}

impl From<tray_icon::Error> for AppError {
    fn from(err: tray_icon::Error) -> Self {
        Self::MenuError(err.to_string())
    }
}

type Result<T> = std::result::Result<T, AppError>;

// 配置结构
#[derive(Serialize, Deserialize, Clone)]
struct Config {
    save_path: PathBuf,
    buffer_duration: Duration,
    voice_enhance: bool,
    voice_gain: f32,
    #[serde(skip)]
    modified: bool,
}

impl Default for Config {
    fn default() -> Self {
        Self {
            save_path: dirs::audio_dir().unwrap_or_default(),
            buffer_duration: Duration::from_secs(1800),
            voice_enhance: true,
            voice_gain: 1.5,
            modified: false,
        }
    }
}

impl Config {
    fn load() -> Result<Self> {
        let config_path = Self::config_path()?;
        if config_path.exists() {
            let content = std::fs::read_to_string(config_path)?;
            Ok(serde_json::from_str(&content)?)
        } else {
            Ok(Self::default())
        }
    }

    fn save(&self) -> Result<()> {
        if self.modified {
            let config_path = Self::config_path()?;
            if let Some(dir) = config_path.parent() {
                std::fs::create_dir_all(dir)?;
            }
            let content = serde_json::to_string_pretty(self)?;
            std::fs::write(config_path, content)?;
        }
        Ok(())
    }

    fn config_path() -> Result<PathBuf> {
        dirs::config_dir()
            .map(|p| p.join("recording-retrospective").join("config.json"))
            .ok_or_else(|| AppError::SystemError("无法获取配置目录".into()))
    }

    fn set_voice_enhance(&mut self, enable: bool) {
        self.voice_enhance = enable;
        self.modified = true;
    }
}

// 音频缓冲区结构
struct AudioBuffer {
    producer: Producer<i16, Arc<HeapRb<i16>>>,
    consumer: Arc<Mutex<Consumer<i16, Arc<HeapRb<i16>>>>>,
}

impl AudioBuffer {
    fn new(duration_secs: u64, sample_rate: u32) -> Self {
        let capacity = (duration_secs as u64 * sample_rate as u64) as usize;
        let max_duration = capacity as f32 / sample_rate as f32;
        
        info!("创建环形缓冲区:");
        info!("- 容量: {} 样本", capacity);
        info!("- 采样率: {} Hz", sample_rate);
        info!("- 最大时长: {:.2} 秒 ({:.2} 分钟)", max_duration, max_duration / 60.0);
        info!("- 内存占用: {:.2} MB", (capacity * std::mem::size_of::<i16>()) as f32 / (1024.0 * 1024.0));
        
        let buffer = HeapRb::new(capacity);
        let (producer, consumer) = buffer.split();
        
        Self {
            producer: producer.into(),
            consumer: Arc::new(Mutex::new(consumer.into())),
        }
    }

    fn reinitialize(&mut self, duration_secs: u64, sample_rate: u32) {
        info!("重新初始化音频缓冲区...");
        let new_buffer = AudioBuffer::new(duration_secs, sample_rate);
        self.producer = new_buffer.producer;
        self.consumer = new_buffer.consumer;
        info!("音频缓冲区重新初始化完成");
    }

    fn clip_sample(sample: i16, gain: f32) -> i16 {
        let scaled = (sample as f32 * gain) as i32;
        scaled.clamp(i16::MIN as i32, i16::MAX as i32) as i16
    }

    fn get_samples_copy(&self, duration_secs: Option<u64>, sample_rate: u32) -> Vec<i16> {
        let consumer = self.consumer.lock();
        let available = consumer.len();
        
        let samples_needed = if let Some(duration) = duration_secs {
            ((duration as u32 * sample_rate) as usize).min(available)
        } else {
            available
        };
        
        let mut samples = Vec::with_capacity(samples_needed);
        
        if samples_needed < available {
            samples.extend(consumer.iter().copied().skip(available - samples_needed));
        } else {
            samples.extend(consumer.iter().copied());
        }
        
        samples
    }

    fn normalize_samples(&self, samples: &mut [i16]) {
        if let Some(max) = samples.iter().map(|s| s.abs() as i32).reduce(i32::max) {
            if max > i16::MAX as i32 {
                let scale = (i16::MAX as f32) / (max as f32);
                samples.iter_mut().for_each(|s| *s = (*s as f32 * scale) as i16);
            }
        }
    }

    fn log_memory_usage(&self) {
        let sample_count = self.consumer.lock().len();
        let memory_usage = (sample_count * std::mem::size_of::<i16>()) / 1024; // KB
        info!("当前音频缓冲区使用内存: {}KB, 样本数: {}", memory_usage, sample_count);
    }
}

// 带通滤波器实现
struct BandpassFilter {
    low_pass: LowPassFilter,
    high_pass: HighPassFilter,
}

struct LowPassFilter {
    prev_output: f32,
    alpha: f32,
}

struct HighPassFilter {
    prev_input: f32,
    prev_output: f32,
    alpha: f32,
}

impl LowPassFilter {
    fn new(cutoff_freq: f32, sample_rate: f32) -> Self {
        let rc = 1.0 / (2.0 * std::f32::consts::PI * cutoff_freq);
        let dt = 1.0 / sample_rate;
        let alpha = dt / (rc + dt);
        Self {
            prev_output: 0.0,
            alpha,
        }
    }

    fn process(&mut self, input: i16) -> f32 {
        let input_f32 = input as f32 / i16::MAX as f32;
        self.prev_output = self.prev_output + self.alpha * (input_f32 - self.prev_output);
        self.prev_output
    }
}

impl HighPassFilter {
    fn new(cutoff_freq: f32, sample_rate: f32) -> Self {
        let rc = 1.0 / (2.0 * std::f32::consts::PI * cutoff_freq);
        let dt = 1.0 / sample_rate;
        let alpha = rc / (rc + dt);
        Self {
            prev_input: 0.0,
            prev_output: 0.0,
            alpha,
        }
    }

    fn process(&mut self, input: f32) -> f32 {
        let output = self.alpha * (self.prev_output + input - self.prev_input);
        self.prev_input = input;
        self.prev_output = output;
        output
    }
}

impl BandpassFilter {
    fn new(low_freq: f32, high_freq: f32, sample_rate: f32) -> Self {
        Self {
            low_pass: LowPassFilter::new(high_freq, sample_rate),
            high_pass: HighPassFilter::new(low_freq, sample_rate),
        }
    }

    fn process(&mut self, input: i16) -> i16 {
        let high_passed = self.high_pass.process(self.low_pass.process(input));
        (high_passed * i16::MAX as f32) as i16
    }
}

// 应用状态结构
struct AppState {
    config: Config,
    audio_buffer: AudioBuffer,
}

impl AppState {
    fn new(config: Config) -> Result<Self> {
        let audio_buffer = AudioBuffer::new(1800, 44100);
        
        Ok(Self {
            config,
            audio_buffer,
        })
    }

    fn update_sample_rate(&mut self, sample_rate: u32) {
        self.audio_buffer.reinitialize(1800, sample_rate);
    }
}

// 为 AppState 添加 Drop 实现
impl Drop for AppState {
    fn drop(&mut self) {
        info!("开始清理资源");
        if let Err(e) = self.config.save() {
            error!("保存配置失败: {}", e);
        }
        info!("资源清理完成");
    }
}

// 音频处理函数
fn start_recording_buffer(state: Arc<Mutex<AppState>>) -> Result<(cpal::Stream, u32)> {
    info!("开始初始化录音设备...");
    
    let host = cpal::default_host();
    info!("成功获取默认音频主机");

    let device = host.default_input_device()
        .ok_or_else(|| AppError::AudioError("无法找到录音设备".into()))?;
    info!("成功获取默认音设备");

    let config = device.default_input_config()?;
    let sample_rate = config.sample_rate().0;
    info!("成功获取设备配置：{:?}", config);
    info!("实际采样率: {}", sample_rate);

    {
        let mut state = state.lock();
        state.update_sample_rate(sample_rate);
    }

    let channels = config.channels() as usize;
    info!("采样率: {}, 声道数: {}", sample_rate, channels);

    let filter = Arc::new(Mutex::new(BandpassFilter::new(200.0, 4000.0, sample_rate as f32)));

    let stream = device.build_input_stream(
        &config.into(),
        move |data: &[f32], _: &cpal::InputCallbackInfo| {
            if let Some(mut state) = state.try_lock() {
                let mut filter = filter.lock();
                let mut processed_samples = Vec::with_capacity(data.len() / channels);
                
                for frame in data.chunks(channels) {
                    // 计算所有声道的平均值
                    let sample_avg = if channels > 1 {
                        frame.iter().sum::<f32>() / channels as f32
                    } else {
                        frame[0]
                    };
                    
                    let sample_i16 = (sample_avg * i16::MAX as f32) as i16;
                    
                    let processed_sample = if state.config.voice_enhance {
                        let filtered = filter.process(sample_i16);
                        AudioBuffer::clip_sample(filtered, state.config.voice_gain)
                    } else {
                        sample_i16
                    };
                    
                    processed_samples.push(processed_sample);
                }
                
                // 批量写入处理后的样本
                for &sample in &processed_samples {
                    if state.audio_buffer.producer.is_full() {
                        let mut consumer = state.audio_buffer.consumer.lock();
                        let _ = consumer.pop();
                    }
                    let _ = state.audio_buffer.producer.push(sample);
                }
            }
        },
        |err| error!("录音错误: {}", err),
        None,
    )?;

    stream.play()?;
    info!("成功启动录音流");
    
    Ok((stream, sample_rate))
}

// 事件枚举
#[derive(Debug)]
enum AppEvent {
    SaveBuffer5Min,
    SaveBuffer10Min,
    SaveBuffer30Min,
    ToggleVoiceEnhance,
    OpenStartup,
    Quit,
}

// 保存录音函数
fn save_recent_recording(state: Arc<Mutex<AppState>>, sample_rate: u32, duration_secs: u64) -> Result<()> {
    {
        let state = state.lock();
        state.audio_buffer.log_memory_usage();
    }
    info!("开始保存录音... 采样率: {}, 预期时长: {}秒", sample_rate, duration_secs);
    
    let mut samples = {
        let state = state.lock();
        state.audio_buffer.get_samples_copy(Some(duration_secs), sample_rate)
    };
    
    let actual_duration = samples.len() as f32 / sample_rate as f32;
    info!("获取到 {} 个样本，实际时长: {:.2}秒", samples.len(), actual_duration);
    
    state.lock().audio_buffer.normalize_samples(&mut samples);
    
    let file_path = {
        let state = state.lock();
        let timestamp = Local::now().format("%y%m%d.%H%M%S");
        let file_name = format!("{}.wav", timestamp);
        state.config.save_path.join(&file_name)
    };

    info!("准备保存到文件: {}", file_path.display());

    let spec = hound::WavSpec {
        channels: 1,
        sample_rate,
        bits_per_sample: 16,
        sample_format: hound::SampleFormat::Int,
    };

    let mut writer = hound::WavWriter::create(&file_path, spec)?;
    
    for sample in samples {
        writer.write_sample(sample)?;
    }
    
    writer.finalize()?;
    info!("录音保存成功！文件位置: {}", file_path.display());
    
    if let Some(file_path_str) = file_path.to_str() {
        std::process::Command::new("explorer")
            .args(["/select,", file_path_str])
            .spawn()
            .map_err(|e| AppError::SystemError(format!("无法打开文件夹: {}", e)))?;
    }
    
    Ok(())
}

// 完整的菜单创建代码
fn create_tray_menu(menu: Menu) -> Result<(
    TrayIcon, 
    crossbeam_channel::Receiver<AppEvent>
)> {
    info!("开始创建托盘菜单...");
    
    let icon = {
        let icon_bytes = include_bytes!("../assets/icon.ico");
        info!("正在加载图标数据...");
        let icon_dir = ico::IconDir::read(std::io::Cursor::new(icon_bytes))
            .map_err(|e| AppError::IconError(format!("图标解析错误: {}", e)))?;
        
        let entry = icon_dir.entries().iter()
            .find(|e| e.width() == 32 && e.height() == 32)
            .ok_or_else(|| AppError::IconError("找不到合适的图标".to_string()))?;
            
        let icon_image = entry.decode()
            .map_err(|e| AppError::IconError(format!("图标解码错误: {}", e)))?;
        
        info!("图标大小: {}x{}", icon_image.width(), icon_image.height());
        
        Icon::from_rgba(
            icon_image.rgba_data().to_vec(),
            icon_image.width(),
            icon_image.height()
        ).map_err(|e| AppError::IconError(format!("创建图标错误: {}", e)))?
    };
    info!("图标加载成功");

    info!("创建菜单...");
    let (app_tx, app_rx) = crossbeam_channel::unbounded();
    let app_tx_clone = app_tx.clone();
    
    let save_buffer_5min = MenuItem::new("保存5分钟录音", true, None);
    let save_buffer_5min_id = MENU_ID_SAVE_5MIN;
    info!("菜单项 ID - 保存5分钟: {}", save_buffer_5min_id);
    menu.append(&save_buffer_5min);
    
    let save_buffer_10min = MenuItem::new("保存10分钟录音", true, None);
    let save_buffer_10min_id = MENU_ID_SAVE_10MIN;
    info!("菜单项 ID - 保存10分钟: {}", save_buffer_10min_id);
    menu.append(&save_buffer_10min);
    
    let save_buffer_30min = MenuItem::new("保存30分钟录音", true, None);
    let save_buffer_30min_id = MENU_ID_SAVE_30MIN;
    info!("菜单项 ID - 保存30分钟: {}", save_buffer_30min_id);
    menu.append(&save_buffer_30min);
    
    let open_startup = MenuItem::new("打开自启动目录", true, None);
    let open_startup_id = MENU_ID_OPEN_STARTUP;
    info!("菜单项 ID - 打开自启动目录: {}", open_startup_id);
    menu.append(&open_startup);
    
    let enhance_voice = CheckMenuItem::new("人声增强", true, true, None);
    let enhance_voice_id = MENU_ID_ENHANCE_VOICE;
    info!("菜单项 ID - 人声增强: {}", enhance_voice_id);
    menu.append(&enhance_voice);
    
    let quit = MenuItem::new("退出", true, None);
    let quit_id = MENU_ID_QUIT;
    info!("菜单项 ID - 退出: {}", quit_id);
    menu.append(&quit);

    info!("构建托盘图标...");
    let tray_icon = TrayIconBuilder::new()
        .with_icon(icon)
        .with_menu(Box::new(menu.clone()))
        .with_tooltip("录音回溯")
        .build()?;

    let menu_event_receiver = MenuEvent::receiver();
    std::thread::spawn(move || {
        while let Ok(event) = menu_event_receiver.recv() {
            let event_id = event.id;
            if let Err(e) = app_tx_clone.send(match event_id {
                MENU_ID_SAVE_5MIN => AppEvent::SaveBuffer5Min,
                MENU_ID_SAVE_10MIN => AppEvent::SaveBuffer10Min,
                MENU_ID_SAVE_30MIN => AppEvent::SaveBuffer30Min,
                MENU_ID_ENHANCE_VOICE => AppEvent::ToggleVoiceEnhance,
                MENU_ID_OPEN_STARTUP => AppEvent::OpenStartup,
                MENU_ID_QUIT => AppEvent::Quit,
                _ => continue,
            }) {
                error!("发送应用事件失败: {}", e);
                break;
            }
        }
    });

    Ok((tray_icon, app_rx))
}

// 修改主函数
fn main() -> Result<()> {
    env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info"))
        .format(|buf, record| {
            use std::io::Write;
            let timestamp = Local::now().format("%Y-%m-%d %H:%M:%S");
            writeln!(
                buf,
                "{} {} [{}] {}",
                timestamp,
                record.level(),
                record.target(),
                record.args()
            )
        })
        .init();

    info!("程序启动");

    let event_loop = EventLoopBuilder::<AppEvent>::with_user_event().build();
    let event_loop_proxy = event_loop.create_proxy();

    let menu = Menu::new();
    let (tray_icon, app_rx) = create_tray_menu(menu)?;
    let _tray_icon = Arc::new(Mutex::new(tray_icon));

    let config = Config::load()?;
    let state = Arc::new(Mutex::new(AppState::new(config)?));

    let (stream, sample_rate) = start_recording_buffer(Arc::clone(&state))?;
    let _stream = Arc::new(stream);

    let event_proxy = event_loop_proxy.clone();
    std::thread::spawn(move || {
        while let Ok(event) = app_rx.recv() {
            if let Err(e) = event_proxy.send_event(event) {
                error!("转发事件到事件循环失败: {}", e);
                break;
            }
        }
    });

    event_loop.run(move |event, _, control_flow| {
        *control_flow = winit::event_loop::ControlFlow::Wait;

        match event {
            winit::event::Event::UserEvent(event) => {
                info!("收到用户事件: {:?}", event);
                match event {
                    AppEvent::SaveBuffer5Min => {
                        let state_clone = Arc::clone(&state);
                        std::thread::spawn(move || {
                            if let Err(e) = save_recent_recording(state_clone, sample_rate, 300) {
                                error!("保存录音失败: {}", e);
                            }
                        });
                    }
                    AppEvent::SaveBuffer10Min => {
                        let state_clone = Arc::clone(&state);
                        std::thread::spawn(move || {
                            if let Err(e) = save_recent_recording(state_clone, sample_rate, 600) {
                                error!("保存录音失败: {}", e);
                            }
                        });
                    }
                    AppEvent::SaveBuffer30Min => {
                        let state_clone = Arc::clone(&state);
                        std::thread::spawn(move || {
                            if let Err(e) = save_recent_recording(state_clone, sample_rate, 1800) {
                                error!("保存录音失败: {}", e);
                            }
                        });
                    }
                    AppEvent::ToggleVoiceEnhance => {
                        let mut state = state.lock();
                        let new_state = !state.config.voice_enhance;
                        state.config.set_voice_enhance(new_state);
                        info!("人声增强: {}", if new_state { "开启" } else { "关闭" });
                    }
                    AppEvent::OpenStartup => {
                        info!("打开自启动目录");
                        if let Some(startup_path) = dirs::config_dir()
                            .map(|p| p.join(r"Microsoft\Windows\Start Menu\Programs\Startup"))
                        {
                            if let Err(e) = std::process::Command::new("explorer")
                                .arg(&startup_path)
                                .spawn() 
                            {
                                error!("无法打开启动文件夹: {}", e);
                            }
                        }
                    }
                    AppEvent::Quit => {
                        info!("退出程序");
                        *control_flow = winit::event_loop::ControlFlow::Exit;
                    }
                }
            }
            winit::event::Event::LoopDestroyed => {
                info!("事件循环结束");
            }
            _ => {}
        }
    })
}

